#include <stdlib.h>
#include <stdio.h>
#include <assert.h>
#include <string.h>

#include <cuda_runtime.h>
#include <math.h>

#include "hb.h"
#include "algos.hpp"

//-----------------------------------------------------------------------------
//-----------------------------------------------------------------------------

Svec::Svec(Config &cfg, int pad, int replicas)
{
    assert(cfg.dims <= 4);
    assert(cfg.cpu_precision == cfg.gpu_precision);
    assert(cfg.cpu_precision == SinglePrecision);

    m_cfg = cfg;

    m_volume = 1;
    for(int i=0; i<m_cfg.dims; ++i) {
        m_volume *= m_cfg.lat[i];
    }
    m_pad = pad;
    m_replicas = replicas;

    size_t element_size;
    size_t vector_size;
    size_t replica_size;
    switch(m_replicas) {
    case 1:
        // pack a 3d vector into a float4
        element_size = sizeof(float);
        vector_size = sizeof(float4);
        replica_size = 1;
        m_stride = 1; // not used
        break;
    case 2:
        // pack 2 consecutive float4
        element_size = sizeof(float);
        vector_size = 2 * sizeof(float4);
        replica_size = 1;
        m_stride = 1; // not used
    case 4:
        // pack a 4 replicas 3d vectors into a float4[3]
        element_size = sizeof(float);
        vector_size = sizeof(float4);
        replica_size = 3;
        m_stride = m_volume + m_pad;
        break;
    default:
        assert(false);
    }
    m_gpu_bytes = replica_size * (m_volume + m_pad) * vector_size;
    m_cpu_bytes = m_gpu_bytes;

    m_cpu_buf = malloc(m_cpu_bytes);
    if(m_cpu_buf == NULL) {
        error("Svec: error allocating CPU memory\n");
    }
    memset(m_cpu_buf, 0, m_cpu_bytes);

    if(cudaMalloc((void**)&m_gpu_buf, m_gpu_bytes) != cudaSuccess) {
        error("Svec: error allocating GPU memory\n");
    }
    if(cudaMemset(m_gpu_buf, 0, m_gpu_bytes) != cudaSuccess) {
        error("Svec: error setting GPU memory\n");
    }
}

//-----------------------------------------------------------------------------

Svec::~Svec()
{
    cudaFree(m_gpu_buf);
    m_gpu_buf = NULL;
    m_gpu_bytes = 0;

    free(m_cpu_buf);
    m_cpu_buf = NULL;
    m_cpu_bytes = 0;
}

//-----------------------------------------------------------------------------

void Svec::cpu_to_gpu()
{
    trace("copying cpu %p -> GPU %p  %d bytes\n", m_cpu_buf, m_gpu_buf, m_gpu_bytes);
    if(cudaMemcpy(m_gpu_buf, m_cpu_buf, m_gpu_bytes, cudaMemcpyHostToDevice) != cudaSuccess) {
        error("%s: error copying\n", __FUNCTION__);
	}
}

//-----------------------------------------------------------------------------

void Svec::gpu_to_cpu()
{
    if(cudaMemcpy(m_cpu_buf, m_gpu_buf, m_cpu_bytes, cudaMemcpyDeviceToHost) != cudaSuccess) {
        error("%s: error copying\n", __FUNCTION__);
	}
}

//-----------------------------------------------------------------------------

void Svec::init_r1(float x, float y, float z)
{
    assert(m_cpu_buf);
    int lx = m_cfg.lat[DIR_X];
    int ly = m_cfg.lat[DIR_Y];
    int lz = m_cfg.lat[DIR_Z];
    float4* vec = (float4*)m_cpu_buf;
    trace("init to (%f,%f,%f)\n", x, y, z);
    for(int ix=0; ix<lx; ++ix)
        for(int iy=0; iy<ly; ++iy)
            for(int iz=0; iz<lz; ++iz) {
                int idx = (iz) + (iy)*lz + (ix)*(lz*ly);
                vec[idx] = make_float4(x,y,z,0.0);
            }
}

//-----------------------------------------------------------------------------

void Svec::init_r4(float x, float y, float z)
{
    assert(m_cpu_buf);
    int lx = m_cfg.lat[DIR_X];
    int ly = m_cfg.lat[DIR_Y];
    int lz = m_cfg.lat[DIR_Z];
    float4* vec = (float4*)m_cpu_buf;
    for(int ix=0; ix<lx; ++ix)
        for(int iy=0; iy<ly; ++iy)
            for(int iz=0; iz<lz; ++iz) {
                int idx = (iz) + (iy)*lz + (ix)*(lz*ly);
                vec[idx+0*m_stride] = make_float4(x,x,x,x);
                vec[idx+1*m_stride] = make_float4(y,y,y,y);
                vec[idx+2*m_stride] = make_float4(z,z,z,z);
            }
}

//-----------------------------------------------------------------------------

void Svec::init_r2(float x, float y, float z)
{
    assert(m_cpu_buf);
    int lx = m_cfg.lat[DIR_X];
    int ly = m_cfg.lat[DIR_Y];
    int lz = m_cfg.lat[DIR_Z];
    float4* vec = (float4*)m_cpu_buf;
    for(int ix=0; ix<lx; ++ix)
        for(int iy=0; iy<ly; ++iy)
            for(int iz=0; iz<lz; ++iz) {
                int idx = (iz) + (iy)*lz + (ix)*(lz*ly);
                vec[2*idx+0] = make_float4(x,y,z,0.0);
                vec[2*idx+1] = make_float4(x,y,z,0.0);
            }
}

//-----------------------------------------------------------------------------

void Svec::init(const vec &v)
{
    init(v.x, v.y, v.z);
}

//-----------------------------------------------------------------------------

void Svec::init(float x, float y, float z)
{
    switch(m_replicas) {
    case 1: init_r1(x,y,z); break;
    case 2: init_r2(x,y,z); break;
    case 4: init_r4(x,y,z); break;
    default: error("invalid replica\n");
    }
}

//-----------------------------------------------------------------------------

static const float norm_eps = 1.0e-4;

bool Svec::check_norm() const
{
    bool ret = true;
    assert(1 == m_replicas);

    const int lx = m_cfg.lat[DIR_X];
    const int ly = m_cfg.lat[DIR_Y];
    const int lz = m_cfg.lat[DIR_Z];

    for(int ix=0; ix<lx; ++ix)
        for(int iy=0; iy<ly; ++iy)
            for(int iz=0; iz<lz; ++iz) {
                float4 f = at(ix, iy, iz);
                float norm = std::sqrt(f.x*f.x + f.y*f.y + f.z*f.z);
                if(fabs(norm-1.0f) > norm_eps) {
                    trace("[%d,%d,%d]={%f,%f,%f} norm=%f\n", ix, iy, iz, f.x, f.y, f.z, norm);
                    ret |= false;
                }
            }
    return ret;
}

//-----------------------------------------------------------------------------

float Svec::calc_alignment(const vec &dir) const
{
    assert(1 == m_replicas);

    const int lx = m_cfg.lat[DIR_X];
    const int ly = m_cfg.lat[DIR_Y];
    const int lz = m_cfg.lat[DIR_Z];

    float ret = 0.0f;
    for(int ix=0; ix<lx; ++ix)
        for(int iy=0; iy<ly; ++iy)
            for(int iz=0; iz<lz; ++iz) {
                float4 f = at(ix, iy, iz);
                float scal_prod = f.x*f.x + f.y*f.y + f.z*f.z;
                ret += scal_prod;
                // if(fabs(scal_prod-1.0f) > norm_eps) {
                //     info("svec[%d,%d,%d]={%f,%f,%f} scal_prod=%f\n", ix, iy, iz, f.x, f.y, f.z, scal_prod);
                // }
            }
    return ret / (lx * ly * lz);
}

//-----------------------------------------------------------------------------

void Svec::dump() const
{
    assert(1 == m_replicas);

    const int lx = m_cfg.lat[DIR_X];
    const int ly = m_cfg.lat[DIR_Y];
    const int lz = m_cfg.lat[DIR_Z];

    float ret = 0.0f;
    for(int ix=0; ix<lx; ++ix)
        for(int iy=0; iy<ly; ++iy)
            for(int iz=0; iz<lz; ++iz) {
                float4 f = at(ix, iy, iz);
                info("[%d,%d,%d]={%f,%f,%f,%f}\n", ix, iy, iz, f.x, f.y, f.z, f.w);
            }
}

//-----------------------------------------------------------------------------
/*
 * Local variables:
 *  mode: c++
 *  c-indent-level: 4
 *  c-basic-offset: 4
 *  tab-width: 4
 *  indent-tabs-mode: nil
 * End:
 */
